#!/bin/bash

# Copyright 2019 Mobvoi Inc. All Rights Reserved.
. ./path.sh || exit 1;

# Use this to control how many gpu you use, It's 1-gpu training if you specify
# just 1gpu, otherwise it's is multiple gpu training based on DDP in pytorch
export CUDA_VISIBLE_DEVICES="0,1,2,3"
# The NCCL_SOCKET_IFNAME variable specifies which IP interface to use for nccl
# communication. More details can be found in
# https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html
# export NCCL_SOCKET_IFNAME=ens4f1
export NCCL_DEBUG=INFO
export CUDA_LAUNCH_BLOCKING=1
stage=5 # start from 0 if you need to start from data preparation
stop_stage=5

# The num of machines(nodes) for multi-machine training, 1 is for one machine.
# NFS is required if num_nodes > 1.
num_nodes=1

# The rank of each node or machine, which ranges from 0 to `num_nodes - 1`.
# You should set the node_rank=0 on the first machine, set the node_rank=1
# on the second machine, and so on.
node_rank=0
# The aishell dataset location, please change this to your own path
# make sure of using absolute path. DO-NOT-USE relatvie path!
data=/export/data/asr-data/OpenSLR/33/
data_url=www.openslr.org/resources/33

nj=16
dict=data/dict/lang_char.txt

# data_type can be `raw` or `shard`. Typically, raw is used for small dataset,
# `shard` is used for large dataset which is over 1k hours, and `shard` is
# faster on reading data and training.
data_type=raw
num_utts_per_shard=1000

train_set=AISHELL1
dev_set=AISHELL_dev
# Optional train_config
# 1. conf/train_transformer.yaml: Standard transformer
# 2. conf/train_conformer.yaml: Standard conformer
# 3. conf/train_unified_conformer.yaml: Unified dynamic chunk causal conformer
# 4. conf/train_unified_transformer.yaml: Unified dynamic chunk transformer
# 5. conf/train_u2++_conformer.yaml: U2++ conformer
# 6. conf/train_u2++_transformer.yaml: U2++ transformer
train_config=conf/tuning/train_asr_whisper_full.yaml
cmvn=false
dir=exp/whisper_large_v0
# dir=exp/whisper_large_aishell1_deepspeed_mn
checkpoint=
num_workers=0
prefetch=0

# use average_checkpoint will get better result
average_checkpoint=false
decode_checkpoint=final.pt
decode_checkpoint=/home/work_nfs5_ssd/yzli/workspace/wenet/examples/aishell/s0/exp/whisper_large_v0/final.pt
# decode_checkpoint=-2.pt
average_num=3
decode_modes="recognize_whisper"
#test_sets="aishell1 aishell2 aishell4 test_net test_meeting Test_Ali_far SPEECHIO_ASR_ZH00003 SPEECHIO_ASR_ZH00002 SPEECHIO_ASR_ZH00001 Test_Ali_near SPEECHIO_ASR_ZH00000"
# test_sets="test/aishell2"
test_sets="test_aishell_8k test/aishell" 
# test_sets="test_ali_short"

if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
  # Test model, please specify the model you want to test by --checkpoint
  if [ ${average_checkpoint} == true ]; then
    decode_checkpoint=avg_${average_num}.pt
    echo "do model average and final checkpoint is $decode_checkpoint"
    python wenet/bin/average_model.py \
      --dst_model $dir/$decode_checkpoint \
      --src_path $dir  \
      --num ${average_num} \
      --val_best
  fi
  # Please specify decoding_chunk_size for unified streaming and
  # non-streaming model. The default value is -1, which is full chunk
  # for non-streaming inference.
  decoding_chunk_size=
  ctc_weight=0.3
  reverse_weight=0.5
  for test_set in ${test_sets}; do
  {
    test_dir=$dir/test_${decode_modes}_${decode_checkpoint}/$test_set
    mkdir -p $test_dir
    python wenet/bin/recognize.py --gpu 0 \
      --mode $decode_modes \
      --config $dir/train.yaml \
      --data_type $data_type \
      --test_data data/$test_set/data.list \
      --checkpoint $decode_checkpoint \
      --beam_size 1 \
      --batch_size 1 \
      --penalty 0.0 \
      --dict $dict \
      --ctc_weight $ctc_weight \
      --reverse_weight $reverse_weight \
      --result_file $test_dir/text \
      ${decoding_chunk_size:+--decoding_chunk_size $decoding_chunk_size}
    python tools/compute-wer.py --char=1 --v=1 \
      data/test/$test_set/text $test_dir/text > $test_dir/wer
  }
  done
  wait
fi
